package copilot

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gofrs/uuid"
	"github.com/tidwall/gjson"
	"github.com/tidwall/sjson"
)

// CodeCompletions 处理代码补全请求。
// 它从客户端接收请求，根据配置修改请求，
// 然后将其发送到代码补全服务，并将响应以流式传输回客户端。
func CodeCompletions(c *gin.Context) {
	ctx := c.Request.Context()

	// 为跟踪生成唯一的请求ID。
	requestID := uuid.Must(uuid.NewV4()).String()
	c.Header("x-github-request-id", requestID)

	// 添加防抖时间以防止过多请求。
	debounceTime, _ := strconv.Atoi(os.Getenv("COPILOT_DEBOUNCE"))
	time.Sleep(time.Duration(debounceTime) * time.Millisecond)

	// 检查客户端是否已取消请求。
	if ctx.Err() != nil {
		abortCodex(c, http.StatusRequestTimeout)
		return
	}

	// 读取请求体。
	body, err := io.ReadAll(c.Request.Body)
	if nil != err {
		abortCodex(c, http.StatusBadRequest)
		return
	}

	// 将响应内容类型设置为 event-stream。
	c.Header("Content-Type", "text/event-stream")
	// 从环境变量中获取代码补全服务类型。
	codexServiceType := os.Getenv("CODEX_SERVICE_TYPE")
	// 为代码补全服务构造请求体。
	body = ConstructRequestBody(body, codexServiceType)

	// 创建一个新的到代码补全服务的HTTP请求。
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, os.Getenv("CODEX_API_BASE"), io.NopCloser(bytes.NewBuffer(body)))
	if nil != err {
		abortCodex(c, http.StatusInternalServerError)
		return
	}

	// 设置请求头。
	req.Header.Set("Content-Type", "application/json")

	// 从环境变量中获取API密钥。
	apiKeys := strings.Split(os.Getenv("CODEX_API_KEY"), ",")

	// 检查是否有任何有效的API密钥。
	if len(apiKeys) == 0 || (len(apiKeys) == 1 && apiKeys[0] == "") {
		abortCodex(c, http.StatusInternalServerError)
		return
	}

	// 选择一个随机的API密钥。
	randGen := rand.New(rand.NewSource(time.Now().UnixNano()))
	selectedKey := strings.TrimSpace(apiKeys[randGen.Intn(len(apiKeys))])

	// 检查所选密钥是否为空。
	if selectedKey == "" {
		abortCodex(c, http.StatusInternalServerError)
		return
	}

	// 使用所选API密钥设置Authorization头。
	req.Header.Set("Authorization", "Bearer "+selectedKey)
	// 创建一个带超时的新HTTP客户端。
	httpClientTimeout, _ := time.ParseDuration(os.Getenv("HTTP_CLIENT_TIMEOUT") + "s")
	client := &http.Client{
		Timeout: httpClientTimeout,
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, // 允许自签名证书的不安全连接。
		},
	}
	// 输出请求地址
	// log.Println("Request URL: ", req.URL)
	// 输出请求参数
	// log.Println("Request Body: ", string(body))
	// 将请求发送到代码补全服务。
	resp, err := client.Do(req)
	if nil != err {
		if errors.Is(err, context.Canceled) {
			abortCodex(c, http.StatusRequestTimeout)
			return
		}

		log.Println("请求补全失败:", err.Error())
		abortCodex(c, http.StatusInternalServerError)
		return
	}
	defer CloseIO(resp.Body)

	// 检查响应状态码。
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		log.Println("请求补全失败:", string(body))

		abortCodex(c, resp.StatusCode)
		return
	}

	c.Status(resp.StatusCode)
	// 处理Ollama服务的流式响应。
	if codexServiceType == "ollama" {
		reader := bufio.NewReader(resp.Body)
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				if err == io.EOF {
					break
				}
				break
			}

			if strings.TrimSpace(line) == "" {
				continue
			}

			// 从流中解析JSON行。
			lineJson := gjson.Parse(line)
			uuid := uuid.Must(uuid.NewV4()).String()
			done := lineJson.Get("done").Bool()
			doneReason := lineJson.Get("done_reason").Str
			response := lineJson.Get("response").Str
			timestamp := time.Now().Unix()
			// 构造 choice 对象。
			choice := map[string]interface{}{
				"text":          response,
				"index":         0,
				"logprobs":      nil,
				"finish_reason": doneReason,
			}
			choices := []map[string]interface{}{choice}
			// 构造要发送给客户端的数据。
			constructLineData := map[string]interface{}{
				"id":                 uuid,
				"choices":            choices,
				"created":            timestamp,
				"model":              lineJson.Get("model").Str,
				"system_fingerprint": "fp_1c141eb703",
				"object":             "text_completion",
			}

			// 如果补全完成，则添加使用情况统计信息。
			if done && strings.Contains(doneReason, "stop") {
				usage := map[string]interface{}{
					"prompt_tokens":            lineJson.Get("prompt_eval_count").Int(),
					"completion_tokens":        lineJson.Get("eval_count").Int(),
					"total_tokens":             lineJson.Get("prompt_eval_count").Int(),
					"prompt_cache_hit_tokens":  lineJson.Get("prompt_eval_count").Int(),
					"prompt_cache_miss_tokens": lineJson.Get("eval_count").Int(),
				}
				constructLineData["usage"] = usage
			}

			// 将修改后的数据重新编码为JSON。
			modifiedJSON, err := json.Marshal(constructLineData)
			if err != nil {
				continue
			}

			// 将修改后的数据发送给客户端。
			_, _ = c.Writer.WriteString("data: " + string(modifiedJSON) + "\n\n")
			c.Writer.Flush()
		}

		// 发送[DONE]消息以指示流的结束。
		_, _ = c.Writer.WriteString("data: [DONE]\n\n")
		c.Writer.Flush()
		return
	}

	// 对于其他服务，只需代理响应体。
	_, _ = io.Copy(c.Writer, resp.Body)
}

// ConstructRequestBody 根据服务类型和环境变量修改请求体。
// 它设置模型，启用流式传输，删除不必要的字段并应用长度限制。
func ConstructRequestBody(body []byte, codexServiceType string) []byte {
	// 从环境变量设置模型名称。
	envCodexModel := os.Getenv("CODEX_API_MODEL_NAME")
	// 从环境变量设置模型服务商类型
	envProviderType := os.Getenv("PROVIDER_TYPE")
	body, _ = sjson.SetBytes(body, "model", envCodexModel)
	// 强制流式输出。
	body, _ = sjson.SetBytes(body, "stream", true)
	// 删除不必要的字段。
	body, _ = sjson.DeleteBytes(body, "extra")
	body, _ = sjson.DeleteBytes(body, "nwo")
	body, _ = sjson.DeleteBytes(body, "logprobs")

	// 对提示和后缀应用长度限制。
	body = applyPromptLengthLimit(body)

	// 从环境变量设置温度。
	temperature, _ := strconv.ParseFloat(os.Getenv("CODEX_TEMPERATURE"), 64)
	if temperature != -1 {
		body, _ = sjson.SetBytes(body, "temperature", temperature)
	}

	// 从环境变量设置 max_tokens，并设置上限。
	codeMaxTokens, _ := strconv.Atoi(os.Getenv("CODEX_MAX_TOKENS"))
	if int(gjson.GetBytes(body, "max_tokens").Int()) > codeMaxTokens {
		body, _ = sjson.SetBytes(body, "max_tokens", codeMaxTokens)
	}

	// 将补全数量限制为1。
	if gjson.GetBytes(body, "n").Int() > 1 {
		body, _ = sjson.SetBytes(body, "n", 1)
	}

	// 根据模型名称构造请求体。
	// 对于 Stable Code 和 Codegemma 模型: https://ollama.com/library/stable-code, https://ollama.com/library/codegemma
	if strings.Contains(envCodexModel, "stable-code") || strings.Contains(envCodexModel, "codegemma") {
		return constructWithStableCodeModel(body)
	}

	// 对于 CodeLlama 模型: https://ollama.com/library/codellama
	if strings.Contains(envCodexModel, "codellama") {
		return constructWithCodeLlamaModel(body)
	}

	// 对于 Qwen Coder Turbo 模型: https://help.aliyun.com/zh/model-studio/user-guide/qwen-coder
	if strings.Contains(envCodexModel, "qwen-coder-turbo") {
		return constructWithQwenCoderTurboModel(body)
	}

	// 对于 new-api和gpt5的模型
	if strings.Contains(envCodexModel, "gpt-5") || strings.Contains(envProviderType, "new-api") {
		log.Println("GPT5!!!!!!!!!!!!!!!!!!")

		return constructWithGpt5Model(body)
	}

	// 对于支持Ollama FIM的模型，如Deepseek Coder v2: https://ollama.com/library/deepseek-coder-v2
	if codexServiceType == "ollama" {
		return constructWithOllamaModel(body, codeMaxTokens)
	}

	return body
}

// applyPromptLengthLimit 对提示和后缀应用长度限制。
func applyPromptLengthLimit(body []byte) []byte {
	// 从环境变量中获取提示长度限制。
	envLimitPrompt := os.Getenv("CODEX_LIMIT_PROMPT")
	limitPrompt, err := strconv.Atoi(envLimitPrompt)
	if err != nil || limitPrompt <= 0 {
		return body // 如果值无效或未设置，则不限制。
	}

	body = limitPromptLength(body, limitPrompt)
	body = limitSuffixLength(body, limitPrompt)

	return body
}

// limitPromptLength 通过取最后 `limitRows` 行来限制提示的长度。
func limitPromptLength(body []byte, limitRows int) []byte {
	prompt := gjson.GetBytes(body, "prompt")
	if !prompt.Exists() {
		return body
	}

	rows := strings.Split(prompt.Str, "\n")
	if len(rows) <= limitRows {
		return body
	}

	// 保留提示的后半部分。
	newPrompt := strings.Join(rows[len(rows)-limitRows:], "\n")
	body, _ = sjson.SetBytes(body, "prompt", newPrompt)

	return body
}

// limitSuffixLength 通过取前 `limitRows` 行来限制后缀的长度。
func limitSuffixLength(body []byte, limitRows int) []byte {
	suffix := gjson.GetBytes(body, "suffix")
	if !suffix.Exists() {
		return body
	}

	rows := strings.Split(suffix.Str, "\n")
	if len(rows) <= limitRows {
		return body
	}

	// 保留后缀的前半部分。
	newSuffix := strings.Join(rows[:limitRows], "\n")
	body, _ = sjson.SetBytes(body, "suffix", newSuffix)

	return body
}

// constructWithCodeLlamaModel 为CodeLlama模型构造请求体。
// 它用特定的标签包裹提示和后缀。
func constructWithCodeLlamaModel(body []byte) []byte {
	body, _ = sjson.SetBytes(body, "stream", true)
	suffix := gjson.GetBytes(body, "suffix")
	prompt := gjson.GetBytes(body, "prompt")
	content := fmt.Sprintf("<PRE> %s <SUF> %s <MID>", prompt, suffix)

	return constructWithChatModel(body, content)
}

// constructWithStableCodeModel 为StableCode模型构造请求体。
// 它用FIM（Fill-in-the-Middle）标签包裹提示和后缀。
func constructWithStableCodeModel(body []byte) []byte {
	body, _ = sjson.SetBytes(body, "stream", true)
	suffix := gjson.GetBytes(body, "suffix")
	prompt := gjson.GetBytes(body, "prompt")
	content := fmt.Sprintf("<fim_prefix>%s<fim_suffix>%s<fim_middle>", prompt, suffix)

	return constructWithChatModel(body, content)
}

// constructWithChatModel 为基于聊天的模型构造请求体。
// 它创建一个带有用户内容的 `messages` 数组。
func constructWithChatModel(body []byte, content string) []byte {
	body, _ = sjson.SetBytes(body, "stream", true)
	// 创建一个新的JSON对象并将其添加到body中。
	messages := []map[string]string{
		{
			"role":    "user",
			"content": content,
		},
	}

	body, _ = sjson.SetBytes(body, "messages", messages)

	// 对尖括号进行反转义。
	jsonStr := string(body)
	jsonStr = strings.ReplaceAll(jsonStr, "\\u003c", "<")
	jsonStr = strings.ReplaceAll(jsonStr, "\\u003e", ">")
	return []byte(jsonStr)
}

// constructWithQwenCoderTurboModel 为QwenCoderTurbo模型构造请求体。
// 它创建一个包含系统、用户和助手角色的 `messages` 数组。
func constructWithQwenCoderTurboModel(body []byte) []byte {
	if gjson.GetBytes(body, "n").Int() > 1 {
		body, _ = sjson.SetBytes(body, "n", 1)
	}
	suffix := gjson.GetBytes(body, "suffix")
	prompt := gjson.GetBytes(body, "prompt")
	codeLanguage := gjson.GetBytes(body, "extra.language")

	messages := []map[string]interface{}{
		{
			"role":    "system",
			"content": "你是一名精通" + codeLanguage.Str + "编程的专家，非常擅长理解和续写代码。",
		},
		{
			"role": "user",
			"content": "Combined with subsequent code snippets, help me complete the code:\n\n" +
				"Code subsequent content:\n```" + codeLanguage.Str + "\n" + suffix.Str + "```\n\n" +
				"Remember:\n" +
				"- Do not generate content outside of the code.\n" +
				"- Do not directly fill in all the code content, the maximum number of lines of code should not exceed 5 lines.\n" +
				"- Answer must refer to the code suffix content, do not exceed the boundary, otherwise repeated code will occur.\n" +
				"- If you don't know how to answer, just reply with an empty string.",
		},
		{
			"role":    "assistant",
			"content": prompt.Str,
			"partial": true,
		},
	}

	body, _ = sjson.SetBytes(body, "messages", messages)
	body, _ = sjson.DeleteBytes(body, "prompt")
	return body
}

// constructWithGpt5Model 为Gpt5模型构造请求体。
// 它创建一个包含系统、用户和助手角色的 `messages` 数组。
func constructWithGpt5Model(body []byte) []byte {
	if gjson.GetBytes(body, "n").Int() > 1 {
		body, _ = sjson.SetBytes(body, "n", 1)
	}
	prompt := gjson.GetBytes(body, "prompt")
	codeLanguage := gjson.GetBytes(body, "extra.language")

	messages := []map[string]interface{}{
		{
			"role":    "system",
			"content": "你是一名精通" + codeLanguage.Str + "编程的专家，非常擅长理解和续写代码。",
		},
		{
			"role": "user",
			"content": "Combined with subsequent code snippets, help me complete the code:\n\n" +
				"Code subsequent content:\n```" + codeLanguage.Str + "\n" + "```\n\n" +
				"Remember:\n" +
				"- Do not generate content outside of the code.\n" +
				"- Do not directly fill in all the code content, the maximum number of lines of code should not exceed 5 lines.\n" +
				"- Answer must refer to the code suffix content, do not exceed the boundary, otherwise repeated code will occur.\n" +
				"- If you don't know how to answer, just reply with an empty string.",
		},
		{
			"role":    "assistant",
			"content": prompt.Str,
			"partial": true,
		},
	}

	body, _ = sjson.SetBytes(body, "messages", messages)
	body, _ = sjson.DeleteBytes(body, "prompt")
	body, _ = sjson.DeleteBytes(body, "suffix")
	body, _ = sjson.DeleteBytes(body, "stop")

	body, _ = sjson.DeleteBytes(body, "max_tokens")
	body, _ = sjson.DeleteBytes(body, "max_output_tokens")

	return body
}

// constructWithOllamaModel 为Ollama模型构造请求体。
// 它设置了各种选项，如温度、停止序列和流式传输。
func constructWithOllamaModel(body []byte, codeMaxTokens int) []byte {
	body, _ = sjson.SetBytes(body, "options.temperature", 0)
	// 处理 'stop' 参数。
	stopArray := gjson.GetBytes(body, "stop").Array()
	stopSlice := make([]interface{}, len(stopArray))
	for i, v := range stopArray {
		stopSlice[i] = v.String()
	}
	body, _ = sjson.SetBytes(body, "options.stop", stopSlice)
	body, _ = sjson.SetBytes(body, "stream", true)

	// 设置要预测的令牌数。
	maxTokens := gjson.GetBytes(body, "max_tokens").Int()
	if int(maxTokens) > codeMaxTokens {
		body, _ = sjson.SetBytes(body, "options.num_predict", codeMaxTokens)
	} else {
		body, _ = sjson.SetBytes(body, "options.num_predict", maxTokens)
	}
	return body
}

// abortCodex 中止请求并向客户端发送[DONE]消息。
func abortCodex(c *gin.Context, status int) {
	c.Header("Content-Type", "text/event-stream")
	c.String(status, "data: [DONE]\n\n")
	c.Abort()
}
